{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3f82a8c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import json\n",
    "import time\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import os\n",
    "from torchvision import transforms, models, datasets\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import ssl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0344666c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 是否继续训练模型的参数\n",
    "def set_parameter_requires_grad(model, feature_extracting):\n",
    "    if feature_extracting:\n",
    "        for param in model.parameters():\n",
    "            param.requires_grad = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d511992b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_device() -> torch.device:\n",
    "    # 是否用GPU训练\n",
    "    train_on_gpu = torch.cuda.is_available()\n",
    "\n",
    "    if not train_on_gpu:\n",
    "        print('CUDA is not available.  Training on CPU ...')\n",
    "    else:\n",
    "        print('CUDA is available!  Training on GPU ...')\n",
    "\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    return device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "71178681",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 选择迁移哪个模型\n",
    "def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):\n",
    "    # 选择合适的模型，不同模型的初始化方法稍微有点区别\n",
    "    model_ft = None\n",
    "    input_size = 0\n",
    "\n",
    "    if model_name == \"resnet\":\n",
    "        \"\"\" \n",
    "        Resnet152\n",
    "        \"\"\"\n",
    "#         # pretrained 表示是否需要下载\n",
    "#         model_ft = models.resnet50(pretrained=False)  #   50\n",
    "        ############################################################################################\n",
    "        model_ft = models.vgg16(pretrained = False)\n",
    "        pre=torch.load(r'./checkpoints-vgg16/vgg16-397923af.pth')\n",
    "        model_ft.load_state_dict(pre)\n",
    "        \n",
    "        \n",
    "        set_parameter_requires_grad(model_ft, feature_extract)\n",
    "        num_ftrs = model_ft.fc.in_features\n",
    "        model_ft.fc = nn.Sequential(nn.Linear(num_ftrs, num_classes),nn.LogSoftmax(dim=1))\n",
    "        input_size = 224\n",
    "    else:\n",
    "        print(\"Invalid model name, exiting...\")\n",
    "        exit()\n",
    "\n",
    "    return model_ft, input_size\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ab7aa9c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Exception ignored in: <function tqdm.__del__ at 0x7f1c12575d30>\n",
      "Traceback (most recent call last):\n",
      "  File \"/opt/conda/lib/python3.8/site-packages/tqdm/std.py\", line 1150, in __del__\n",
      "    self.close()\n",
      "  File \"/opt/conda/lib/python3.8/site-packages/tqdm/notebook.py\", line 271, in close\n",
      "    self.sp(bar_style='danger')\n",
      "AttributeError: 'tqdm' object has no attribute 'sp'\n"
     ]
    }
   ],
   "source": [
    "def lab_start():\n",
    "    # 种类数\n",
    "    num_classes = 29\n",
    "    model_ft, input_size = initialize_model(\"resnet\", num_classes, feature_extract=True, use_pretrained=True)\n",
    "    device = get_device()\n",
    "    model_ft = model_ft.to(device)\n",
    "    # 带name的params\n",
    "    params = model_ft.named_parameters()\n",
    "    print(\"Params need to learn:\")\n",
    "    params_need_update = []\n",
    "    for param_name, param in params:\n",
    "        if param.requires_grad:\n",
    "            params_need_update.append(param)\n",
    "            print(param_name)\n",
    "    #\n",
    "    #####################   文件夹          ##############################################\n",
    "    #\n",
    "    data_dir = './pics3.0/'\n",
    "    train_dir = data_dir + '/train'\n",
    "    valid_dir = data_dir + '/valid'\n",
    "\n",
    "    # 数据增强\n",
    "    data_transforms = {\n",
    "        'train': transforms.Compose([transforms.RandomRotation(25),  # 随机旋转，-45到45度之间随机选\n",
    "#                                      transforms.Resize(256), # 重新固定大小 #########\n",
    "#                                      transforms.CenterCrop(224),  # 从中心开始裁剪，只得到一张图片\n",
    "                                     transforms.RandomHorizontalFlip(p=0.5),  # 随机水平翻转 概率为0.5\n",
    "                                     transforms.RandomVerticalFlip(p=0.5),  # 随机垂直翻转\n",
    "                                     transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),\n",
    "                                     # 参数1为亮度，参数2为对比度，参数3为饱和度，参数4为色相\n",
    "                                     transforms.RandomGrayscale(p=0.025),  # 概率转换成灰度率，3通道就是R=G=B\n",
    "                                     transforms.ToTensor(),\n",
    "                                     # 迁移学习，用别人的均值和标准差\n",
    "                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 均值，标准差\n",
    "                                     ]),\n",
    "        'valid': transforms.Compose([transforms.Resize(256),\n",
    "                                     transforms.CenterCrop(224),\n",
    "                                     transforms.ToTensor(),\n",
    "                                     # 预处理必须和训练集一致\n",
    "                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                                     ]),\n",
    "    }\n",
    "################################################################################################\n",
    "    batch_size = 8\n",
    "\n",
    "    # train和valid的图片，做transform之后用字典保存\n",
    "    image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in\n",
    "                      ['train', 'valid']}\n",
    "\n",
    "    print(datasets.ImageFolder(os.path.join(data_dir, 'train'), data_transforms['train']).classes)\n",
    "    # 批量处理，这里都是tensor格式（上面compose）\n",
    "    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True) for x in\n",
    "                   ['train', 'valid']}\n",
    "    # fuc = datasets.ImageFolder(os.path.join(data_dir, 'train'),data_transforms['train']).class_to_idx\n",
    "    # with open('fuc_index.txt', 'w') as file:\n",
    "    #     file.write(str(fuc))\n",
    "    # print(fuc)\n",
    "\n",
    "\n",
    "    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'valid']}\n",
    "    print(\"dataset_sizes below:\")\n",
    "    print( dataset_sizes)\n",
    "    # 样本数据的标签\n",
    "    class_names = image_datasets['train'].classes\n",
    "    print(class_names)\n",
    "    #####################################################################################################\n",
    "    # 优化器设置 \n",
    "    optimizer_ft = optim.Adam(params_need_update, lr=1e-2)\n",
    "    # 学习率衰减\n",
    "    scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.15)  # 学习率每7个epoch衰减成原来的1/10\n",
    "    # 最后一层已经LogSoftmax()了，所以不能nn.CrossEntropyLoss()来计算了，nn.CrossEntropyLoss()相当于logSoftmax()和nn.NLLLoss()整合\n",
    "    criterion = nn.NLLLoss()  # 可以尝试改成交叉熵损失函数0702\n",
    "    \n",
    "    \n",
    "    ########################################################################\n",
    "    filename = \"./checkpoints-vgg16/vgg16-0810-1.pth\"\n",
    "\n",
    "    model_ft, val_acc_history, train_acc_history, valid_losses, train_losses, LRs = wz_model_train(model_ft,\n",
    "                                                                                                   dataloaders,\n",
    "                                                                                                   criterion,\n",
    "                                                                                                   optimizer_ft,\n",
    "                                                                                                   scheduler,\n",
    "                                                                                                   filename,\n",
    "                                                                                                   device)\n",
    "    for param in model_ft.parameters():\n",
    "        param.requires_grad = True\n",
    "\n",
    "    # 再继续训练所有的参数，学习率调小一点 ###################################################################\n",
    "    optimizer = optim.Adam(params_need_update, lr=1e-3)\n",
    "    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.1)\n",
    "\n",
    "    # 损失函数\n",
    "    criterion = nn.NLLLoss()\n",
    "    checkpoint = torch.load(filename)\n",
    "    best_acc = checkpoint['best_acc']\n",
    "    model_ft.load_state_dict(checkpoint['state_dict'])\n",
    "    optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "\n",
    "    model_ft, val_acc_history, train_acc_history, valid_losses, train_losses, LRs = wz_model_train(model_ft,\n",
    "                                                                                                   dataloaders,\n",
    "                                                                                                   criterion,\n",
    "                                                                                                   optimizer,\n",
    "                                                                                                   scheduler, filename,\n",
    "                                                                                                   device, #########################################\n",
    "                                                                                                   num_epochs=40)  # 再训练5个纪元\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1703d715",
   "metadata": {},
   "outputs": [],
   "source": [
    "############################################################################################\n",
    "def wz_model_train(model, dataloaders, criterion, optimizer, scheduler, filename: str, device: torch.device,\n",
    "                   num_epochs=100, is_inception=False):  # 原来为 10 ，现在为25 ,num_epochs 为 训练的大阶段，共几个阶段\n",
    "    start_time = time.time()\n",
    "    best_acc = 0\n",
    "    best_model_weights = copy.deepcopy(model.state_dict())\n",
    "    model.to(device)\n",
    "    # 保存损失和准确率数据\n",
    "    val_acc_history = []\n",
    "    train_acc_history = []\n",
    "    train_losses = []\n",
    "    valid_losses = []\n",
    "    # 记录每个epoch的learningRate\n",
    "    LRs = [optimizer.param_groups[0]['lr']]\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        print('Epoch {}/{}'.format(epoch, num_epochs - 1))\n",
    "        print('-' * 10)\n",
    "\n",
    "        # 训练和验证\n",
    "        for phase in ['train', 'valid']:\n",
    "            if phase == 'train':\n",
    "                model.train()  # 训练\n",
    "            else:\n",
    "                model.eval()  # 验证\n",
    "\n",
    "            running_loss = 0.0\n",
    "            running_corrects = 0\n",
    "            for inputs, labels in dataloaders[phase]:\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "                # 清零\n",
    "                optimizer.zero_grad()\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    # inception会有辅助输出，损失函数为一个线性模型\n",
    "                    if is_inception and phase == 'train':\n",
    "                        outputs, aux_outputs = model(inputs)\n",
    "                        loss1 = criterion(outputs, labels)\n",
    "                        loss2 = criterion(aux_outputs, labels)\n",
    "                        loss = loss1 + 0.4 * loss2\n",
    "                    else:  # resnet执行的是这里\n",
    "                        outputs = model(inputs)\n",
    "                        loss = criterion(outputs, labels)\n",
    "\n",
    "                    _, preds = torch.max(outputs, 1)\n",
    "                    # 训练阶段更新权重\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "                # 计算损失\n",
    "                # loss计算默认都是取mean，计算批量的loss时，要乘以loss的数量\n",
    "                # 所以这里计算的是一个epoch里所有样本的loss和正确数量\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                running_corrects += torch.sum(preds == labels.data)\n",
    "            # 完整一次的loss均值和准确率\n",
    "            epoch_loss = running_loss / len(dataloaders[phase].dataset)\n",
    "            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)\n",
    "            # 一个epoch里train和valid分别花的时间和loss和准确度\n",
    "            time_elapsed = time.time() - start_time\n",
    "            print('Time elapsed {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n",
    "            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))\n",
    "            # 得到最好那次的模型\n",
    "            if phase == 'valid' and epoch_acc > best_acc:  # 最大如何比较 0624:就是单纯数值比较0629\n",
    "                best_acc = epoch_acc\n",
    "                best_model_weights = copy.deepcopy(model.state_dict())\n",
    "                state = {\n",
    "                    'state_dict': model.state_dict(),\n",
    "                    'best_acc': best_acc,\n",
    "                    'optimizer': optimizer.state_dict(),\n",
    "                }\n",
    "                torch.save(state, filename)\n",
    "            if phase == 'valid':\n",
    "                val_acc_history.append(epoch_acc)\n",
    "                valid_losses.append(epoch_loss)\n",
    "                scheduler.step()\n",
    "            if phase == 'train':\n",
    "                train_acc_history.append(epoch_acc)\n",
    "                train_losses.append(epoch_loss)\n",
    "        print('Optimizer learning rate : {:.7f}'.format(optimizer.param_groups[0]['lr']))\n",
    "        LRs.append(optimizer.param_groups[0]['lr'])\n",
    "\n",
    "    time_elapsed = time.time() - start_time\n",
    "    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n",
    "    print('Best val Acc: {:4f}'.format(best_acc))\n",
    "    # 训练完后用最好的一次当做模型最终的结果\n",
    "    model.load_state_dict(best_model_weights)\n",
    "    return model, val_acc_history, train_acc_history, valid_losses, train_losses, LRs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bd533d65",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# print(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "bbb7cd8c",
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'VGG' object has no attribute 'fc'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_542/440188911.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mssl\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_create_default_https_context\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mssl\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_create_unverified_context\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mlab_start\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/tmp/ipykernel_542/1456166982.py\u001b[0m in \u001b[0;36mlab_start\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m     \u001b[0;31m# 种类数\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m     \u001b[0mnum_classes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m29\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m     \u001b[0mmodel_ft\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0minitialize_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"resnet\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_classes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeature_extract\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0muse_pretrained\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m     \u001b[0mdevice\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_device\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m     \u001b[0mmodel_ft\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel_ft\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_542/362961264.py\u001b[0m in \u001b[0;36minitialize_model\u001b[0;34m(model_name, num_classes, feature_extract, use_pretrained)\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m         \u001b[0mset_parameter_requires_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel_ft\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeature_extract\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 20\u001b[0;31m         \u001b[0mnum_ftrs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel_ft\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0min_features\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     21\u001b[0m         \u001b[0mmodel_ft\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSequential\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mLinear\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnum_ftrs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_classes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mLogSoftmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m         \u001b[0minput_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m224\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__getattr__\u001b[0;34m(self, name)\u001b[0m\n\u001b[1;32m   1133\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mmodules\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1134\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mmodules\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1135\u001b[0;31m         raise AttributeError(\"'{}' object has no attribute '{}'\".format(\n\u001b[0m\u001b[1;32m   1136\u001b[0m             type(self).__name__, name))\n\u001b[1;32m   1137\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'VGG' object has no attribute 'fc'"
     ]
    }
   ],
   "source": [
    "ssl._create_default_https_context = ssl._create_unverified_context\n",
    "lab_start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1f51cf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
